import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from Res18 import ResNet18
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np

device = torch.device("cpu")

# 数据加载与预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5071, 0.4865, 0.4409),
                         (0.2673, 0.2564, 0.2761)),
])

trainset = torchvision.datasets.CIFAR100(root='./data', train=True,
                                         download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR100(root='./data', train=False,
                                        download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64,
                                         shuffle=False, num_workers=2)

# 初始化模型
model = ResNet18(num_classes=100).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1,
                      momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)

train_losses, train_accuracies, test_accuracies = [], [], []

def train(epoch):
    model.train()
    running_loss, correct, total = 0.0, 0, 0
    loop = tqdm(trainloader, desc=f"Epoch {epoch} [Train]")
    for inputs, targets in loop:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        loop.set_postfix(loss=running_loss / total,
                         acc=100. * correct / total)

    train_losses.append(running_loss / total)
    train_accuracies.append(100. * correct / total)

def test(epoch):
    model.eval()
    correct_top1, correct_top5, total = 0, 0, 0
    all_preds, all_targets = [], []

    with torch.no_grad():
        for inputs, targets in testloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            _, pred_top1 = outputs.topk(1, dim=1)
            _, pred_top5 = outputs.topk(5, dim=1)

            correct_top1 += pred_top1.eq(targets.view(-1, 1)).sum().item()
            correct_top5 += pred_top5.eq(targets.view(-1, 1).expand_as(pred_top5)).sum().item()
            total += targets.size(0)

            all_preds.extend(pred_top1.view(-1).cpu().numpy())
            all_targets.extend(targets.cpu().numpy())

    acc1 = 100. * correct_top1 / total
    acc5 = 100. * correct_top5 / total
    test_accuracies.append(acc1)
    print(f"Epoch {epoch} [Test] Top-1 Acc: {acc1:.2f}%, Top-5 Acc: {acc5:.2f}%")

    if epoch == num_epochs:
        cm = confusion_matrix(all_targets, all_preds)
        disp = ConfusionMatrixDisplay(cm)
        disp.plot(cmap='Blues', xticks_rotation='vertical')
        plt.title(f"Confusion Matrix (Epoch {epoch})")
        plt.savefig("confusion_matrix.png")
        plt.close()

def plot_metrics():
    epochs = range(1, num_epochs + 1)
    plt.figure(figsize=(10, 4))
    plt.subplot(1, 2, 1)
    plt.plot(epochs, train_losses, label='Train Loss')
    plt.xlabel("Epoch"); plt.ylabel("Loss")
    plt.grid(True); plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(epochs, train_accuracies, label='Train Acc')
    plt.plot(epochs, test_accuracies, label='Test Top-1 Acc')
    plt.xlabel("Epoch"); plt.ylabel("Accuracy (%)")
    plt.grid(True); plt.legend()

    plt.tight_layout()
    plt.savefig("training_metrics.png")
    plt.close()

# 主训练流程
if __name__ == '__main__':
    num_epochs = 5
    for epoch in range(1, num_epochs + 1):
        train(epoch)
        test(epoch)
        scheduler.step()

    torch.save(model.state_dict(), "resnet18_cifar100_cpu.pth")
    plot_metrics()

